from absl import app, flags

import os
import ujson as json
import glob
import logging

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def plot_model(cmodel, thres):
  plt.plot(cmodel['probs'], cmodel['quantiles'], lw=0.5)
  plt.minorticks_on()
  plt.grid(lw=0.5, which='major')
  plt.grid(lw=0.2, which='minor')
  plt.hlines(thres, xmin=0, xmax=1, lw=0.2, color='g')


def main(_):
  lin_alphas = glob.glob("profile/%s/linear_alpha.json" % flags.FLAGS.profile)
  os.makedirs("pic_profile", exist_ok=True)

  plt.rcParams['figure.figsize'] = 13, 6
  plt.rcParams['font.size'] = 5
  plt.rcParams['legend.fontsize'] = 5
  plt.rcParams['xtick.labelsize'] = 5
  plt.rcParams['ytick.labelsize'] = 5

  for lin_alpha in lin_alphas:
    lasetting = json.load(open(lin_alpha))
    try:
      cmodel = json.load(open(lin_alpha.replace("linear_alpha", "chosen_models")))
    except Exception:
      continue
    nplots = len(lasetting['model_threshold'].items())
    cnt = 1
    for mk, thres in lasetting['model_threshold'].items():
      for side in ['buy_model', 'sell_model']:
        for each_model in cmodel[side]:
          if each_model['model_id'] == mk:
            plt.subplot(2, int(nplots / 2), cnt)
            plot_model(each_model, thres)
            plt.title(mk)
            cnt += 1
    plt.tight_layout()
    plt.savefig("pic_profile/%s.png" % lin_alpha.split("/")[1])
    plt.close()


if __name__ == '__main__':
  flags.DEFINE_string('profile', '*', 'profile')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
